package com.suqiu.qqrobot.config;

import cn.hutool.core.collection.CollectionUtil;
import com.suqiu.qqrobot.config.properties.OpenAiProperties;
import com.unfbx.chatgpt.entity.chat.Message;
import lombok.Data;
import lombok.EqualsAndHashCode;

import java.util.ArrayList;
import java.util.HashMap;

/**
 * openai 聊天上下文
 */
@EqualsAndHashCode(callSuper = true)
@Data
public class OpenAiChatContext extends HashMap<String, ArrayList<Message>> {

    private OpenAiProperties openAiProperties;

    public OpenAiChatContext(OpenAiProperties openAiProperties) {
        this.openAiProperties = openAiProperties;
    }

    @Override
    public ArrayList<Message> put(String sessionId, ArrayList<Message> value) {
        if (sessionId == null || sessionId.isEmpty()) {
            throw new RuntimeException("sessionId 格式应为 string");
        }
        ArrayList<Message> messages = super.get(sessionId);
        if (CollectionUtil.isEmpty(messages)) {
            messages = new ArrayList<>();
        }
        // 上下文长度 > 配置的数量，则保留最新一条
        if (messages.size() > openAiProperties.getChatGpt().getMaxContext()) {
            Message lastMessage = messages.get(messages.size() - 1);
            messages.clear();
            messages.add(lastMessage);
        }
        return super.put(sessionId, value);
    }
}
